import tiktoken

"""
pip install tiktoken
pip install torch
"""

tokenizer = tiktoken.get_encoding('gpt2')

with open('../data/the-verdict.txt', 'r', encoding='utf-8') as f:
    raw_text = f.read()

enc_text = tokenizer.encode(raw_text)
print(len(enc_text))
# 5145

# 从数据集中移除前50个词元以便演示
enc_sample = enc_text[50:]

# 创建下一单词预测任务的输入-目标对的一种简单且直观的方法是定义两个变量：x和y。
# 变量x用于存储输入的词元，变量y则用于存储由x的每个输入词元右移一个位置所得的目标词元：

context_size = 4
x = enc_sample[:context_size]
y = enc_sample[1:context_size + 1]
print(f'x:  {x}')
print(f'y:      {y}')

# 通过处理输入及其相应的目标（将输入右移一个位置），可以创建图2-12中的下一单词预测任务
for i in range(1, context_size + 1):
    context = enc_sample[:i]
    desired = enc_sample[i]
    print(context, "---->", desired)
# [290] ----> 4920
# [290, 4920] ----> 2241
# [290, 4920, 2241] ----> 287
# [290, 4920, 2241, 287] ----> 257


# 为了更直观地展示这一过程，让我们重用前面的代码，但这一次将词元ID转换回文本形式
for i in range(1, context_size + 1):
    context = enc_sample[:i]
    desired = enc_sample[i]
    print(tokenizer.decode(context), "---->", tokenizer.decode([desired]))


# and ---->  established
# and established - --->  himself
# and established himself - ---> in
# and established himself in ---->  a
